# src/image_processing/key_sensitivity_analysis.py

import numpy as np
import os
from src.image_processing.image_handler import ImageData

def calculate_npcr(img1, img2):
    """
    计算像素数变化率 (NPCR)
    :param img1: 第一张图像数组
    :param img2: 第二张图像数组
    :return: 各通道的 NPCR 值
    """
    if img1.shape != img2.shape:
        raise ValueError("两张图像的形状必须相同！")

    if len(img1.shape) == 2:  # 灰度图像
        diff = img1 != img2
        npcr = np.sum(diff) / (img1.shape[0] * img1.shape[1]) * 100
        return [npcr]  # 返回单通道的 NPCR
    else:  # 多通道图像（如 RGB 或 RGBA）
        npcr_values = []
        for i in range(img1.shape[2]):  # 遍历每个通道
            diff = img1[:, :, i] != img2[:, :, i]
            npcr_values.append(np.sum(diff) / (img1.shape[0] * img1.shape[1]) * 100)
        return npcr_values  # 返回各通道的 NPCR

def calculate_uaci(img1, img2):
    """
    计算平均变化强度 (UACI)
    :param img1: 第一张图像数组
    :param img2: 第二张图像数组
    :return: 各通道的 UACI 值
    """
    if img1.shape != img2.shape:
        raise ValueError("两张图像的形状必须相同！")

    if len(img1.shape) == 2:  # 灰度图像
        diff = np.abs(img1.astype(np.int16) - img2.astype(np.int16))
        uaci = np.sum(diff) / 255 / (img1.shape[0] * img1.shape[1]) * 100
        return [uaci]  # 返回单通道的 UACI
    else:  # 多通道图像（如 RGB 或 RGBA）
        uaci_values = []
        for i in range(img1.shape[2]):  # 遍历每个通道
            diff = np.abs(img1[:, :, i].astype(np.int16) - img2[:, :, i].astype(np.int16))
            uaci_values.append(np.sum(diff) / 255 / (img1.shape[0] * img1.shape[1]) * 100)
        return uaci_values  # 返回各通道的 UACI

def analyze_key_sensitivity(image_data: ImageData, original_key, modified_key, algorithm_name, save_dir, title):
    """
    分析密钥的敏感度
    :param image_data: ImageData 对象
    :param original_key: 原始密钥
    :param modified_key: 修改后的密钥
    :param algorithm_name: 算法名称（如 "RC4" 或 "LFSR"）
    :param save_dir: 保存加密图像的目录
    :param title: 图像标题（包含算法信息）
    :return: 各通道的 NPCR 和 UACI 值，以及图像模式
    """
    mode = image_data.get_mode()
    # 确保算法名称是小写的
    algorithm_name = algorithm_name.lower()

    # 生成加密图像的文件名，基于传入的 title
    original_output_path = os.path.join(save_dir, f'{title}_original_encrypted.png')
    modified_output_path = os.path.join(save_dir, f'{title}_modified_encrypted.png')

    # 使用 ImageData 类加密图像
    original_encrypted_data = image_data.encrypt(original_key, algorithm_name, original_output_path)
    modified_encrypted_data = image_data.encrypt(modified_key, algorithm_name, modified_output_path)

    # 计算 NPCR 和 UACI
    img1 = original_encrypted_data.get_image_array()
    img2 = modified_encrypted_data.get_image_array()
    npcr_values = calculate_npcr(img1, img2)
    uaci_values = calculate_uaci(img1, img2)

    # 返回 NPCR、UACI 和图像模式
    return npcr_values, uaci_values, mode
